"""
Description: UR5 motion simulation server.

Author: Jin Liu
Date: 2025/05/27
"""

import numpy as np
from typing import Union
from pathlib import Path
from qtpy import QtCore, QtGui, QtWidgets
from ezgl.items import GLAxisItem, PointLight, GLGridItem, GLURDFItem, GLLinePlotItem, GLInstancedMeshItem, GLMeshItem, axis_mesh
from ezgl import GLViewWidget, tb, Matrix4x4, MeshData
from ezrobot.utils.logger import ColoredLogger
from ezrobot.utils.rpc import RpcServerThread
from ezrobot.motion import IKSolver

BASE_DIR = Path(__file__).resolve().parent


class RobotServer(GLViewWidget):
    """
    Robot Server Widget.
    """
    def __init__(
            self,
            parent=None,
            urdf_path: Union[str, Path] = BASE_DIR / "../assets/ur5/ur5.urdf",
            ee_name: str = "ee_link",
            base_name: str = "world",
            **kwargs
        ):
        """
        Initialization.
        """
        super().__init__(parent=parent, **kwargs)

        self.logger = ColoredLogger.get_logger("RobotServer")
        self.logger.setLevel("DEBUG")
        self.hand_t_ee = Matrix4x4()
        self.urdf_path = urdf_path
        self.ee_name = ee_name
        self.init_ui()
        self.init_server()
        self.init_ik(base_name, ee_name)

        timer = QtCore.QTimer(self)
        timer.timeout.connect(self.on_timeout)
        timer.start(20)

        self.rpc_server.start()

    def init_ik(self, base_name, ee_name):
        """
        初始化逆运动学求解器.
        """
        last_joint = None
        for joint in self.model_item.get_joints(movable=False):
            if joint.child.name == ee_name:
                last_joint = joint.name
                break                
        self._ik_solover = IKSolver(urdf_path=self.urdf_path, base_elements=[base_name], last_joint=last_joint)
        
    def init_server(self):
        """
        初始化服务器, 注册函数
        """
        self.rpc_server = RpcServerThread(quiet=True)
        self.rpc_server.register_function(self.set_joint, "set_joint")
        self.rpc_server.register_function(self.get_joint, "get_joint")
        self.rpc_server.register_function(self.get_tf, "get_tf")
        self.rpc_server.register_function(self.get_pose, "get_pose")
        self.rpc_server.register_function(self.set_ee, "set_ee")
        self.rpc_server.register_function(self.get_ee, "get_ee")
        self.rpc_server.register_function(self.add_pose_axis_at_ee, "add_pose_axis_at_ee")
        self.rpc_server.register_function(self.add_pose_axis, "add_pose_axis")
        self.rpc_server.register_function(self.set_pose_axis, "set_pose_axis")
        self.rpc_server.register_function(self.add_arrow, "add_arrow")
        self.rpc_server.register_function(self.set_arrow, "set_arrow")
        self.rpc_server.register_function(self.clear_axes, "clear_axes")
        self.rpc_server.register_function(self.toggle_traj, "toggle_traj")
        self.rpc_server.register_function(self.set_cartesian, "set_cartesian")
        self.rpc_server.register_function(self.set_named_axis, "set_named_axis")
        self.logger.info(f"Server started at port: {self.rpc_server.port}")

    def init_ui(self):
        """
        初始化 UI
        """
        # self.resize(1200, 900)
        self.camera.set_vector6d((0.0, 0.0, -3.8, 0, -70, -90))
        self.ax_item = GLAxisItem(fix_to_corner=True)

        # -- lights
        self.light = PointLight(pos=(15, 7, 10), ambient=(1, 1, 1), diffuse=(1, 1, 1),
                               visible=False, directional=True, render_shadow=True)
        self.light.setShadow(-3, 3, -3, 3, -3, 3, bias=0.05)
        # -- grid
        self.grid_item = GLGridItem(
            size=(5, 5), spacing=(0.25, 0.25), lineWidth=1, color=[1,0.92,0.76],
            lights=self.light
        ).rotate(90, 1, 0, 0)
        # self.grid_item.setVisible(0)
        # -- urdf
        self.model_item = GLURDFItem(
            self.urdf_path,
            lights=self.light,
            glOptions="translucent_cull"
        )
        self.model_item.print_links()

        # -- End Effector
        vert_axis, ind_axis, color_axis = axis_mesh(0.15, 0.004)
        self.end_effector = GLInstancedMeshItem(MeshData.Mesh(vert_axis, ind_axis), pose=[Matrix4x4()], color=color_axis, color_divisor=0)
        self.end_effector.setVisible(False)
        self.end_effector.setTransform(self.get_pose())

        # -- trajectory
        self.trajectory_item = GLLinePlotItem(lineWidth=2, color=(0., 1., 0, 1))
        self.trajectory_item.setVisible(False)
        vert_axis, ind_axis, color_axis = MeshData.axis_mesh(0.015, 0.0015, tip_width=1.5, tip_length=0.2)
        self.traj_axes_item = GLInstancedMeshItem(MeshData.Mesh(vert_axis, ind_axis), color=color_axis, color_divisor=0)

        # -- arrows
        vert_arrow, ind_arrow = MeshData.arrow_mesh(1, 0.003, tip_width=1.5, tip_length=0.2)
        self.arrows = GLInstancedMeshItem(MeshData.Mesh(vert_arrow, ind_arrow), color=(1,0,0,1), color_divisor=1, glOptions="ontop")

        # -- 5 named axes
        vert_axis, ind_axis, color_axis = MeshData.axis_mesh(0.06, 0.004, tip_width=1.2, tip_length=0.1)
        color_axis[:, :3] -= 0.5
        color_transform = [Matrix4x4.fromEulerAngles(30, 10, 0), Matrix4x4.fromEulerAngles(-30, -10, 0), 
                           Matrix4x4.fromEulerAngles(0, 30, 10), Matrix4x4.fromEulerAngles(0, -30, -30), Matrix4x4.fromEulerAngles(10, 10, 30)]
        colors = [tf * color_axis + 0.5 for tf in color_transform]
        self.named_axes = [GLInstancedMeshItem(MeshData.Mesh(vert_axis, ind_axis), pose=Matrix4x4(), color=color, color_divisor=0) for color in colors]
        self.addItems(self.named_axes)

        # add items
        self.addItems([self.ax_item, self.grid_item, self.model_item,
                       self.trajectory_item, self.traj_axes_item, self.end_effector, self.arrows])

        # tool_box
        j_value = self.model_item.get_joints_attr("value")
        j_name = self.model_item.get_joints_attr("name")
        j_limits = np.array(self.model_item.get_joints_attr("limit"))
        links_name = self.model_item.get_links_name()

        with tb.window("control", None, 10, size=(1000, 500), pos=(400, 0), frameless=False):
            with tb.splitter(True, spacing=10, stretchs=(5, 2)):
                tb.add_widget("3d", self)
                with tb.group("control", horizontal=False, collapsible=False, show=False):
                    tb.add_drag_array(
                        "joints",
                        value = j_value,
                        min_val = j_limits[:, 0], max_val = j_limits[:, 1],
                        step=[0.001]*len(j_value), decimals=[3]*len(j_value),
                        format=[name+": %.3f" for name in j_name],
                        callback=self.on_changed, horizontal=False
                    )
                    with tb.group("axis_visible", horizontal=False, collapsible=True):
                        tb.add_checklist("axis_visible", items=links_name, callback=self.on_changed,
                                        exclusive=False, horizontal=False, value=None)
                    with tb.group("trajectory", horizontal=False, collapsible=True):
                        with tb.group("checkboxes", horizontal=True, show=False, spacing=15):
                            tb.add_checkbox("show_trajectory", value=False, callback=self.toggle_traj)
                            tb.add_checkbox("show_ee", value=False, callback=self.end_effector.setVisible)
                        tb.add_button("add_axis_at_ee", callback=self.add_pose_axis_at_ee)
                        tb.add_button("clear_axes", callback=self.clear_axes)
                    tb.add_stretch(1)

    def on_timeout(self):
        self.update()

    def on_changed(self, data):
        """
        toolbox changed callback.

        Parameters:
        - data : 控件返回的数据
        """
        label = self.sender().get_label()
        if label == "joints":
            tag, val = data
            self.model_item.set_joint(tag, val)
            self._update_after_motion()
        elif label == "axis_visible":
            tag, val = data
            self.model_item.set_link(tag, axis_visiable=val)

    def closeEvent(self, a0) -> None:
        tb.clean()
        return super().closeEvent(a0)

    def _update_after_motion(self):
        """
        更新轨迹
        """
        o_t_ee = self.get_pose()
        self.end_effector.setTransform(o_t_ee)
        if self.trajectory_item.visible():
            pt = o_t_ee.toTranslation()
            self.trajectory_item.addData(pt)

    def get_joint(self) -> list:
        """
        读取机器人 7 自由度关节角度 (rad)

        Returns:
        - list[float], 长度为 7 的列表
        """
        joint = [float(q) for q in self.model_item.get_joints_attr("value")]
        return joint[:7]  # exclude the last two joints

    def set_joint(self, values: Union[list, np.ndarray, tuple]):
        """
        设置机器人 7 自由度关节角度 (rad)

        Parameters:
        - values : list[float], 长度为 <=9 的列表, 7 自由度 + 2 个夹爪
        """
        self.model_item.set_joints(values)
        tb.get_widget("joints").set_value_no_signal(values)
        self._update_after_motion()

    def set_cartesian(self, pose: Matrix4x4):
        """
        设置笛卡尔空间末端执行器位姿, 内部先求反解, 然后设置关节角度

        Parameters:
        - pose : Matrix4x4, ee 位姿矩阵
        """
        base_t_hand = pose * self.hand_t_ee.inverse()
        joints_rad = self._ik_solover.inverse_kinematics(base_t_hand, initial_guess=self.get_joint())
        self.set_joint(joints_rad)    
        
    def get_tf(self, link:str) -> Matrix4x4:
        """
        读取机器人指定 link 的位姿

        Parameters:
        - link : str, link name
        """
        return self.model_item.get_link_tf(link)

    def get_pose(self) -> Matrix4x4:
        """
        获取末端执行器位姿
        """
        return self.model_item.get_link_tf(self.ee_name) * self.hand_t_ee

    def get_ee(self) -> Matrix4x4:
        """
        获取末端执行器相对于手腕的位姿
        """
        return self.hand_t_ee

    def set_ee(self, hand_t_ee: Matrix4x4):
        """
        设置末端执行器相对于手腕的位姿

        Parameters:
        - hand_t_ee : Matrix4x4
        """
        self.hand_t_ee = hand_t_ee
        self._update_after_motion()

    def toggle_traj(self, visible: bool):
        """
        显示/隐藏轨迹
        """
        if visible:
            self.trajectory_item.setVisible(True)
        else:
            self.trajectory_item.setVisible(False)
            self.trajectory_item.setData([])

    def add_pose_axis(self, pose: Matrix4x4):
        """
        在指定位置添加一个坐标轴
        """
        self.traj_axes_item.addData(pose=pose)

    def set_pose_axis(self, pose: Matrix4x4):
        """
        设置一个坐标轴的位置
        """
        self.traj_axes_item.setData(pose=pose)

    def add_arrow(self, st_pos: np.ndarray, ed_pos: np.ndarray, color: np.ndarray):
        """
        在指定位置添加一个箭头
        """
        pose = MeshData.direction_matrixs(st_pos, ed_pos, False)[0].T
        pose[:3, 2] *= np.linalg.norm(np.array(ed_pos) - np.array(st_pos))
        self.arrows.addData(pose=Matrix4x4(pose), color=color)

    def set_arrow(self, st_pos: np.ndarray, ed_pos: np.ndarray, color: np.ndarray):
        """
        设置一个箭头
        """
        pose = MeshData.direction_matrixs(st_pos, ed_pos, False)[0].T
        pose[:3, 2] *= np.linalg.norm(np.array(ed_pos) - np.array(st_pos))
        self.arrows.setData(pose=Matrix4x4(pose), color=color)

    def add_pose_axis_at_ee(self):
        """
        在当前 EE 位置添加一个轨迹点可视化
        """
        self.traj_axes_item.addData(pose=self.get_pose())

    def clear_axes(self):
        """
        移除最后所有轨迹点
        """
        self.traj_axes_item.setData([])
        for axis in self.named_axes:
            axis.setVisible(False)

    def set_named_axis(self, index, pose=None):
        """
        设置指定索引的命名坐标轴的位姿
        """
        if pose is None:
            self.named_axes[index].setVisible(False)
        else:
            self.named_axes[index].setVisible(True)
            self.named_axes[index].setTransform(pose)        
        
        

if __name__ == '__main__':
    import sys
    QtWidgets.QApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)
    app = QtWidgets.QApplication(sys.argv)
    win = RobotServer(None)
    app.setStyleSheet("""
        QWidget {
            background-color: white;
            color: black;
        }
    """)
    sys.exit(app.exec_())